package com.star.spring.ai.controller;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.ai.image.ImageResponse;
import org.springframework.ai.zhipuai.ZhiPuAiChatModel;
import org.springframework.ai.zhipuai.ZhiPuAiChatOptions;
import org.springframework.ai.zhipuai.ZhiPuAiImageModel;
import org.springframework.ai.zhipuai.ZhiPuAiImageOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

@Slf4j
@RestController
public class ZhipuaiController {

    @Autowired
    private ZhiPuAiChatModel zhiPuAiChatModel;

    @Autowired
    private ZhiPuAiImageModel zhiPuAiImageModel;

    @GetMapping(value = "/zhipuai/generateStream", produces = "text/html;charset=UTF-8")
    public Flux<String> generateStream(@RequestParam(value = "message", defaultValue = "给我讲一个笑话") String message) {
        Prompt prompt = new Prompt(new UserMessage(message));
        return zhiPuAiChatModel.stream(prompt).map((r) -> r.getResult() != null && r.getResult().getOutput() != null && r.getResult().getOutput().getContent() != null ? r.getResult().getOutput().getContent() : "").filter(StringUtils::hasLength);
    }

    @GetMapping(value = "/zhipuai/functionCalling", produces = "text/html;charset=UTF-8")
    public String functionCalling(@RequestParam(value = "message", defaultValue = "今天的热点新闻都有哪些") String message) {
        Prompt prompt = new Prompt(new UserMessage(message), ZhiPuAiChatOptions.builder().withFunction("hotNews").build());
        return zhiPuAiChatModel.call(prompt).getResult().getOutput().getContent();
    }

    @GetMapping(value = "/zhipuai/generateImage")
    public String generateImage(@RequestParam(value = "message", defaultValue = "A light cream colored mini golden doodle") String message) {
        ImageResponse response = zhiPuAiImageModel.call(new ImagePrompt(message, ZhiPuAiImageOptions.builder().withUser("hdd").build()));
        log.info(response.getResult().getOutput().getUrl());
        return response.getResult().getOutput().getUrl();
    }

}
